from IPython.core.display import Image
from IPython.core.display_functions import display
from langgraph.graph import StateGraph
from langgraph.graph import START,END

def addition(state):
    #注意 这里接受到的是初始化状态
    print(f"init_state: {state}")
    return {"x": state['x'] + 1}

def subtraction(state):
    #注意 这里接受到的上一个节点状态
    print(f"init_state: {state}")
    return {"x": state['x'] - 2}

if __name__ == '__main__':
    # pip install pyppeteer ipython
    builder = StateGraph(dict)
    print(builder)
    builder.add_node("addition", addition)
    builder.add_node("subtraction", subtraction)
    builder.add_edge(START, "addition")
    builder.add_edge("addition", "subtraction")
    builder.add_edge("subtraction", END)
    print(builder.edges)
    print(builder.nodes)
    graph=builder.compile()
    graph_image = graph.get_graph(xray=True).draw_mermaid_png()
    with open("graph.png", "wb") as f:
        f.write(graph_image)
    print("Graph saved to graph.png")
    initial_state={"x":10}
    x=graph.invoke(initial_state)
    print(x)